{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ad6e5ed4-a38d-46a6-8bb5-32d68bd0b9e5",
   "metadata": {},
   "source": [
    "End of week 2 exercise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "8f45fbfa-eaaa-4eb8-841e-83b068b80507",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import json\n",
    "import gradio as gr\n",
    "import base64\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "from pydub import AudioSegment\n",
    "from pydub.playback import play\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "451b72a6-1e6c-476a-8431-1c30c5cd9fb8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "API key found and looks good so far!\n"
     ]
    }
   ],
   "source": [
    "# Initialization\n",
    "\n",
    "load_dotenv(override=True)\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "if openai_api_key:\n",
    "    print(\"API key found and looks good so far!\")\n",
    "else:\n",
    "    print(\"No API key was found!\")\n",
    "\n",
    "MODEL = \"gpt-4o-mini\"\n",
    "openai = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "29fa6a53-4b57-47ea-89a1-640020e603b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = (\n",
    "    \"You are a helpful tennis coach who answers questions about tennis rules, \"\n",
    "    \"players, strategies, training, and equipment.\"\n",
    "    \"Give short, courteous answers, no more than 2 sentence.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e9b255e7-02d8-4350-b5d4-e645d1fc90d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Translation\n",
    "\n",
    "LANG_CODES = {\n",
    "    \"English\": \"en\",\n",
    "    \"Spanish\": \"es\",\n",
    "    \"French\": \"fr\"\n",
    "}\n",
    "\n",
    "def translate_text(text, target_language=\"en\"):\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": f\"You are a translator. Translate the following text to {target_language}\"},\n",
    "        {\"role\": \"user\", \"content\": text}\n",
    "    ]\n",
    "    response = openai.chat.completions.create(model=MODEL, messages=messages)\n",
    "    return response.choices[0].message.content\n",
    "\n",
    "def tennis_info_tool(query):\n",
    "    if \"top\" in query.lower():\n",
    "        return \"Top male players: Djokovic, Nadal, Federer. Top female players: Barty, Sabalenka, Swiatek.\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "b44b147c-bfba-4137-9ecb-d5538f08a46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image\n",
    "def generate_tennis_image(prompt):\n",
    "    image_response = openai.images.generate(\n",
    "        model=\"dall-e-3\",\n",
    "        prompt=f\"Tennis scene: {prompt}, realistic and detailed, vibrant colors\",\n",
    "        size=\"1024x1024\",\n",
    "        n=1,\n",
    "        response_format=\"b64_json\",\n",
    "    )\n",
    "    image_base64 = image_response.data[0].b64_json\n",
    "    image_data = base64.b64decode(image_base64)\n",
    "    return Image.open(BytesIO(image_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "bbfeff3b-0c73-4b2c-a6da-3cac27d8fedd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Audio\n",
    "\n",
    "def talker(message):\n",
    "    response = openai.audio.speech.create(\n",
    "        model=\"tts-1\",\n",
    "        voice=\"onyx\",\n",
    "        input=message\n",
    "    )\n",
    "\n",
    "    audio_stream = BytesIO(response.content)\n",
    "    audio = AudioSegment.from_file(audio_stream, format=\"mp3\")\n",
    "    play(audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa338332-3dfc-4b95-8367-65853a8d2793",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "574d672e-0a75-4af9-b3ad-8dc2dec4e607",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(history, user_message, target_language=\"English\", use_audio=False, generate_image=False):\n",
    "    image = None\n",
    "\n",
    "    if any(keyword in user_message.lower() for keyword in [\"top\", \"players\"]):\n",
    "        reply = tennis_info_tool(user_message)\n",
    "    else:\n",
    "        messages = [{\"role\": \"system\", \"content\": system_message}] + history\n",
    "        response = openai.chat.completions.create(model=MODEL, messages=messages)\n",
    "        reply = response.choices[0].message.content\n",
    "    \n",
    "    if target_language != \"English\":\n",
    "        code = LANG_CODES.get(target_language, \"en\")\n",
    "        reply = translate_text(reply, code)\n",
    "\n",
    "    history.append({\"role\": \"assistant\", \"content\": reply})\n",
    "    \n",
    "    if use_audio:\n",
    "        talker(reply)\n",
    "\n",
    "    if generate_image:\n",
    "        image = generate_tennis_image(reply)\n",
    "    return history, image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaf4f47e-d20b-41f8-94b5-4aef0302731b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "9110343f-0efa-49bc-8d5f-498fd690dd14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7869\n",
      "* To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7869/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Input #0, wav, from '/var/folders/73/0s09hh2n48q7s14tld64q3rh0000gn/T/tmp4hoe_x5n.wav':\n",
      "  Duration: 00:00:06.55, bitrate: 384 kb/s\n",
      "  Stream #0:0: Audio: pcm_s16le ([1][0][0][0] / 0x0001), 24000 Hz, 1 channels, s16, 384 kb/s\n",
      "   6.45 M-A: -0.000 fd=   0 aq=    0KB vq=    0KB sq=    0B "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Input #0, wav, from '/var/folders/73/0s09hh2n48q7s14tld64q3rh0000gn/T/tmp2mxw0wth.wav':\n",
      "  Duration: 00:00:04.61, bitrate: 384 kb/s\n",
      "  Stream #0:0: Audio: pcm_s16le ([1][0][0][0] / 0x0001), 24000 Hz, 1 channels, s16, 384 kb/s\n",
      "   4.48 M-A:  0.000 fd=   0 aq=    0KB vq=    0KB sq=    0B "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Gradio\n",
    "\n",
    "with gr.Blocks() as ui:\n",
    "    with gr.Row():\n",
    "        chatbot = gr.Chatbot(height=500, type=\"messages\")\n",
    "        image_output = gr.Image(height=500)\n",
    "    with gr.Row():\n",
    "        entry = gr.Textbox(label=\"Ask your Tennis Coach:\", placeholder=\"Type and press Enter...\")\n",
    "        target_lang_dropdown = gr.Dropdown(\n",
    "            choices=[\"English\", \"Spanish\", \"French\"],\n",
    "            value=\"English\",\n",
    "            label=\"Translate to:\"\n",
    "        )\n",
    "        audio_toggle = gr.Checkbox(label=\"Play audio\", value=False)\n",
    "        image_toggle = gr.Checkbox(label=\"Generate image for this reply\", value=True)\n",
    "    with gr.Row():\n",
    "        clear = gr.Button(\"Clear\")\n",
    "        \n",
    "    def add_user_message(message, history):\n",
    "        history.append({\"role\": \"user\", \"content\": message})\n",
    "        return \"\", history\n",
    "\n",
    "    def chat_response(history, message, target_language, use_audio, generate_image):\n",
    "        history, image = chat(history, message, target_language, use_audio, generate_image)\n",
    "        return history, image\n",
    "\n",
    "    entry.submit(\n",
    "        add_user_message,\n",
    "        inputs=[entry, chatbot],\n",
    "        outputs=[entry, chatbot]).then(\n",
    "        chat_response,\n",
    "        inputs=[chatbot, entry, target_lang_dropdown, audio_toggle, image_toggle],\n",
    "        outputs=[chatbot, image_output]\n",
    "    )\n",
    "\n",
    "    clear.click(lambda: (None, None, None), inputs=None, outputs=[chatbot, image_output, entry], queue=False)\n",
    "\n",
    "ui.launch(inbrowser=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0359c29-22aa-4156-9afa-8c63c02ca747",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
